import cv2
import numpy as np
import os
from pathlib import Path

class ZikaExtractor:
    def __init__(self):
        self.rows = 5  # 5行
        self.cols = 6  # 每行6个

    def preprocess_image(self, image):
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用高斯模糊减少噪声
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        # 使用Otsu's二值化方法
        _, binary = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        
        # 使用形态学操作增强线条
        kernel = np.ones((3, 3), np.uint8)
        dilated = cv2.dilate(binary, kernel, iterations=1)
        
        return dilated

    def find_lines(self, preprocessed_image):
        height, width = preprocessed_image.shape[:2]
        
        # 计算水平投影
        h_proj = np.sum(preprocessed_image, axis=1) / 255
        
        # 设置阈值
        h_threshold = width / 7  # 设置在1/8到1/6之间
        min_line_length = width // 7  # 最小连续线长度为图片宽度的1/4
        
        # 查找横线位置
        h_lines = []
        for i in range(1, len(h_proj)-1):
            if h_proj[i] > h_threshold:
                # 检查这一行是否存在足够长的连续线段
                row = preprocessed_image[i]
                current_length = 0
                max_length = 0
                
                # 统计这一行中最长的连续线段
                for j in range(width):
                    if row[j] > 0:  # 如果是黑色像素
                        current_length += 1
                        max_length = max(max_length, current_length)
                    else:
                        current_length = 0
                
                # 只有当存在足够长的连续线段时才认为是横线
                if max_length >= min_line_length:
                    # 确保是局部最大值
                    window = 5
                    is_peak = True
                    for j in range(max(0, i-window), min(len(h_proj), i+window+1)):
                        if j != i and h_proj[j] > h_proj[i]:
                            is_peak = False
                            break
                    if is_peak:
                        h_lines.append(i)
        
        # 合并相近的横线
        def merge_lines(lines, distance=15):
            if not lines:
                return []
            lines = sorted(lines)
            merged = []
            current_group = [lines[0]]
            
            for line in lines[1:]:
                if line - current_group[-1] <= distance:
                    current_group.append(line)
                else:
                    merged.append(int(np.mean(current_group)))
                    current_group = [line]
            
            if current_group:
                merged.append(int(np.mean(current_group)))
            
            return merged
        
        h_lines = merge_lines(h_lines)
        
        # 确保横线数量正确（应该有6条）
        if len(h_lines) > 6:
            # 保留最强的6条横线
            h_strengths = [h_proj[i] for i in h_lines]
            strongest_indices = np.argsort(h_strengths)[-6:]
            h_lines = [h_lines[i] for i in sorted(strongest_indices)]
        
        # 计算竖线位置（考虑左右边距）
        margin = width // 70  # 左右各预留边距
        available_width = width - 2 * margin  # 减去左右边距后的可用宽度
        col_width = available_width // 6  # 将可用宽度均分为6列
        
        v_lines = []
        for i in range(7):  # 7条线分隔6列
            x = margin + i * col_width  # 从左边距开始，每次增加一个列宽
            v_lines.append(x)
        
        return h_lines, v_lines

    def detect_zika(self, image, preprocessed):
        height, width = image.shape[:2]
        
        # 查找水平和垂直线
        h_lines, v_lines = self.find_lines(preprocessed)
        
        print(f"Found {len(h_lines)} horizontal lines and {len(v_lines)} vertical lines")
        
        # 确保找到足够的线
        if len(h_lines) < 2 or len(v_lines) < 2:
            print("Warning: Not enough lines detected")
            return []
        
        # 根据线的位置划分区域
        regions = []
        
        # 对每个相邻的线对之间的区域进行检查
        for i in range(len(v_lines) - 1):  # 遍历所有竖线对
            for j in range(len(h_lines) - 1):
                x1 = v_lines[i]
                x2 = v_lines[i + 1]
                y1 = h_lines[j]
                y2 = h_lines[j + 1]
                
                # 计算区域尺寸
                w = x2 - x1
                h = y2 - y1
                
                # 检查区域是否合适（不要太小）
                if w > 50 and h > 50:
                    # 略微缩小裁剪区域，以去除边框
                    x1 = x1 + 5
                    y1 = y1 + 5
                    x2 = x2 - 5
                    y2 = y2 - 5
                    
                    regions.append((x1, y1, x2-x1, y2-y1))
        
        # 按照从左到右，从上到下排序
        regions.sort(key=lambda x: (x[1], x[0]))
        
        return regions

    def extract_zika(self, image_path, output_dir):
        # 创建输出目录
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # 读取图片
        image = cv2.imread(str(image_path))
        if image is None:
            raise ValueError(f"Could not read image: {image_path}")
        
        # 获取图片尺寸
        height, width = image.shape[:2]
        print(f"Image size: {width}x{height}")
        
        # 预处理图片
        preprocessed = self.preprocess_image(image)
        
        # 保存预处理后的图片用于调试
        cv2.imwrite(str(output_dir / "preprocessed.jpg"), preprocessed)
        
        # 检测字卡
        h_lines, v_lines = self.find_lines(preprocessed)  # 获取原始检测的线条
        regions = self.detect_zika(image, preprocessed)
        
        # 在原图上绘制检测到的线条和字卡位置
        debug_image = image.copy()
        
        # 绘制横线（蓝色）
        for y in h_lines:
            cv2.line(debug_image, (0, y), (width, y), (255, 0, 0), 2)
            # 在线条旁边标注序号
            cv2.putText(debug_image, f"H{h_lines.index(y)}", (10, y-5),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)
        
        # 绘制竖线（红色）
        for x in v_lines:
            cv2.line(debug_image, (x, 0), (x, height), (0, 0, 255), 2)
            # 在线条上方标注序号
            cv2.putText(debug_image, f"V{v_lines.index(x)}", (x-10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
        
        # 绘制字卡区域（绿色）
        for i, (x, y, w, h) in enumerate(regions):
            cv2.rectangle(debug_image, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(debug_image, str(i+1), (x+10, y+30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # 保存标注后的图片
        cv2.imwrite(str(output_dir / "detected_lines.jpg"), debug_image)
        
        # 提取并保存每个字卡
        for i, (x, y, w, h) in enumerate(regions):
            zika = image[y:y+h, x:x+w]
            cv2.imwrite(str(output_dir / f"zika_{i+1}.jpg"), zika)
            
        return len(regions)

def main():
    # 创建提取器实例
    extractor = ZikaExtractor()
    
    # 获取当前目录
    current_dir = Path.cwd()
    
    # 设置输入输出路径
    input_dir = current_dir / "input"
    output_base_dir = current_dir / "output"
    
    # 确保输入目录存在
    if not input_dir.exists():
        input_dir.mkdir(parents=True)
        print(f"Created input directory: {input_dir}")
        print("Please put your zika images in the input directory and run the script again.")
        return
    
    # 处理输入目录中的所有图片
    image_files = [f for f in input_dir.glob("*") if f.suffix.lower() in {'.png', '.jpg', '.jpeg', '.webp'}]
    
    if not image_files:
        print("No image files found in the input directory.")
        print("Supported formats: PNG, JPG, JPEG, WEBP")
        return
    
    print(f"Found {len(image_files)} images to process.")
    
    # 处理每个图片
    for image_file in image_files:
        print(f"\nProcessing {image_file.name}...")
        
        # 为每个输入图片创建一个输出目录
        output_dir = output_base_dir / image_file.stem
        
        try:
            num_zikas = extractor.extract_zika(image_file, output_dir)
            print(f"Successfully extracted {num_zikas} zikas to {output_dir}")
        except Exception as e:
            print(f"Error processing {image_file.name}: {str(e)}")

if __name__ == "__main__":
    main()
